# -*- coding: utf-8 -*-
"""
Created on Thu May 11 15:23:51 2023

@author: lenovo
"""

import numpy as np
import pandas as pd
from osgeo import gdal
# import os
# import glob
# import shutil
# from tqdm import tqdm # 运行计时

################################################
def read_img(filename):
    dataset = gdal.Open(filename)  # 打开文件
    if dataset == None:
        print(filename+"文件无法打开")
        return
    im_width = dataset.RasterXSize  # 栅格矩阵的列数
    im_height = dataset.RasterYSize  # 栅格矩阵的行数
    im_bands = dataset.RasterCount #波段数
    im_geotrans = dataset.GetGeoTransform()  # 仿射矩阵
    im_proj = dataset.GetProjection()  # 地图投影信息
    im_data = dataset.ReadAsArray(0, 0, im_width, im_height).astype(np.float32)  # 将数据写成数组，对应栅格矩阵
    im_lon=[im_geotrans[0]+i*im_geotrans[1] for i in range(im_width)]
    im_lat=[im_geotrans[3]+i*im_geotrans[5] for i in range(im_height)]
    
    del dataset  # 关闭对象，文件dataset   
    return im_data,im_width,im_height,im_bands,im_geotrans,im_proj,im_lon,im_lat
#=========================
# 保存tif文件函数
def write_img(im_data, im_geotrans, im_proj, path, nodata=None):
    if 'int8' in im_data.dtype.name:
        datatype = gdal.GDT_Byte
    elif 'int16' in im_data.dtype.name:
        datatype = gdal.GDT_UInt16
    else:
        datatype = gdal.GDT_Float32
    if len(im_data.shape) == 3:
        im_bands, im_height, im_width = im_data.shape
    elif len(im_data.shape) == 2:
        im_data = np.array([im_data])
        im_bands, im_height, im_width = im_data.shape
    # 创建文件
    driver = gdal.GetDriverByName("GTiff")
    dataset = driver.Create(path, int(im_width), int(im_height), int(im_bands), datatype)
    if (dataset != None):
        dataset.SetGeoTransform(im_geotrans)  # 写入仿射变换参数
        dataset.SetProjection(im_proj)  # 写入投影        
    for i in range(im_bands):
        if (nodata != None):
            dataset.GetRasterBand(i + 1).SetNoDataValue(nodata)
        dataset.GetRasterBand(i + 1).WriteArray(im_data[i])
    del dataset

#======================================    


vartype=['水源涵养','土壤保持','洪水调蓄','空气净化','水质净化','固碳价值','释氧价值',
         '气候调节','负氧离子','储量碳价值','水资源价值']


tifpath=r'H:\company\公司项目\江西\上饶广丰区\data\GEP'
# df=pd.read_excel(excelpath)

cols=3601
rows=3601
tot=np.zeros((rows,cols))
for tp in vartype:
    fn=tifpath+'\\'+tp+'_N28E118.tif'
    im_data,im_width,im_height,im_bands,im_geotrans,im_proj,im_lon,im_lat=read_img(fn)
    
    im_data[im_data==-9999]=np.nan
    
    tot=tot+im_data

other=(45.3301+203.31+5.30272146+0.0746766+1.85+87.84+126.29)*1e8/1377/1e6
tot=tot+other
tot[np.isnan(tot)]=-9999
fout=tifpath+'\\总价值_N28E118.tif'
write_img(tot, im_geotrans, im_proj, fout,-9999)    
    
